import pdfplumber
from PIL import Image
import io
import fitz  # PyMuPDF

'''
PDF提取模块
'''
class PDFExtractor:
    def __init__(self, file_path):
        self.file_path = file_path
        self.pdf = None
        
    def open_pdf(self):
        """打开PDF文件"""
        try:
            self.pdf = pdfplumber.open(self.file_path)
            return True
        except Exception as e:
            print(f"Error opening PDF file: {e}")
            return False
            
    def extract_image_by_keyword(self, keyword, page_range=None):
        """根据关键词提取图片"""
        images = []
        
        if page_range is None:
            page_range = range(len(self.pdf.pages))
            
        for page_num in page_range:
            page = self.pdf.pages[page_num]
            text = page.extract_text()
            
            if text and keyword in text:
                # 提取页面中的图片
                for img in page.images:
                    # 将图片数据转换为PIL图像
                    image = Image.open(io.BytesIO(img['stream']))
                    images.append((page_num + 1, image))
                    
        return images
        
    def extract_region_by_keyword(self, keyword, buffer=5):
        """根据关键词提取区域"""
        regions = []
        
        for page in self.pdf.pages:
            text = page.extract_text()
            
            if text and keyword in text:
                # 查找关键词的位置
                words = page.extract_words()
                for word in words:
                    if keyword in word['text']:
                        # 计算区域
                        x0 = max(0, word['x0'] - buffer)
                        y0 = max(0, word['top'] - buffer)
                        x1 = min(page.width, word['x1'] + buffer)
                        y1 = min(page.height, word['bottom'] + buffer)
                        
                        # 提取区域图片
                        image = page.crop((x0, y0, x1, y1)).to_image()
                        regions.append((page.page_number, image))
                        
        return regions
        
    def close(self):
        """关闭PDF文件"""
        if self.pdf:
            self.pdf.close()